from wenet.whisper.convert_whisper_to_wenet_config_and_ckpt import main
import whisper
from gxl_ai_utils.utils import utils_file


def _get_whisper_source_pt(pt_dir: str = './data/whisper', whisper_model: str = 'large-v3'):
    utils_file.makedir_sil(pt_dir)
    whisper.load_model(whisper_model, download_root=pt_dir, device='cpu')


def convert_to_wenet(target_pt_dir: str = './output/', source_pt_dir: str = './data/whisper/',
                    whisper_model: str = 'large-v3'):
    utils_file.makedir_sil(target_pt_dir)
    # _get_whisper_source_pt(source_pt_dir, whisper_model)
    main(target_pt_dir, source_pt_dir, whisper_model)


if __name__ == '__main__':
    convert_to_wenet()
